#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Date : 2024/7/15
# @Author : shiyou pan

"""
NOTE:
"""
from itertools import combinations


def _gen_xs_simple(xs: list[str], n_xs: int | range) -> list[list[str]]:
    """ 指定指标清单，根据指标数量进行指标组合的穷举
    params
    ------
    xs: list[str]
        待组合的指标清单
    n_xs: int | range
        组合的长度

    Returns
    -------
    list[list[str]]
        指标组合的列表
    """
    if isinstance(n_xs, int):
        return [list(x) for x in combinations(xs, n_xs)]
    elif isinstance(n_xs, range):
        list_xs = []
        for n_x in n_xs:
            list_xs += _gen_xs_simple(xs, n_x)
        return list_xs


def _gen_xs_wdrop(xs: list[str], n_xs: int | range, xs_drop: list[str] | str | None = None) -> list[list[str]]:
    """
    指定指标清单，根据指标数量进行指标组合的穷举，其中需要排除某些指标
    params
    ------
    xs: list[str]
        待组合的指标清单
    n_xs: int | range
        组合的长度
    xs_drop: list[str] | str | None
        待排除的指标清单，可以是字符串或列表

    Returns
    -------
    list[list[str]]
        指标组合的列表
    """
    if xs_drop is None:
        return _gen_xs_simple(xs, n_xs)
    elif isinstance(xs_drop, list):
        xs = [x for x in xs if x not in xs_drop]
        return _gen_xs_simple(xs, n_xs)
    elif isinstance(xs_drop, str):
        xs = [x for x in xs if x != xs_drop]
        return _gen_xs_simple(xs, n_xs)
    else:
        raise ValueError('xs或xs_drop参数格式错误')


def _gen_xs_wbase(xs: list[str], n_xs: int | range, xs_base: list[str] | None = None) -> list[list[str]]:
    """ xs的组合，仅聚焦存在的指标清单 """
    if xs_base is None:
        return _gen_xs_simple(xs, n_xs)
    elif isinstance(xs_base, list):
        xs = [x for x in xs if x in xs_base]
        return _gen_xs_simple(xs, n_xs)
    else:
        raise ValueError('xs或xs_base参数格式错误')


def _gen_xs_wbd(xs: list[str],
                n_xs: int | range,
                xs_base: list[str] | None = None,
                xs_drop: list[str] | str | None = None) -> list[list[str]]:
    """ xs的组合，聚焦存在的指标清单，排除xs_drop """
    if xs_base is None and xs_drop is None:
        return _gen_xs_simple(xs, n_xs)
    elif xs_base is not None and xs_drop is None:
        return _gen_xs_wbase(xs, n_xs, xs_base)
    elif xs_base is None and xs_drop is not None:
        return _gen_xs_wdrop(xs, n_xs, xs_drop)
    else:
        xs_base = [x for x in xs_base if x not in xs_drop]
        return _gen_xs_wbase(xs, n_xs, xs_base)


def _gen_xslist_wbd(xs: list[list[str]],
                    n_xs: int | range | None = None,
                    xs_base: list[str] | None = None,
                    xs_drop: list[str] | str | None = None) -> list[list[str]]:
    """ xs的组合，聚焦存在的指标清单，排除xs_drop """
    if n_xs is None:
        pass
    elif isinstance(n_xs, int):
        xs = [x for x in xs if len(x) == n_xs]
    elif isinstance(n_xs, range):
        xs = [x for x in xs if len(x) in n_xs]
    else:
        raise ValueError('n_xs参数格式错误')

    if isinstance(xs_drop, str):
        xs_drop = [xs_drop]
    else:
        pass

    if xs_base is None and xs_drop is None:
        return xs
    elif xs_base is not None and xs_drop is None:
        return [x for x in xs if set(x) <= set(xs_base)]
    elif xs_base is None and xs_drop is not None:
        return [x for x in xs if not set(xs_drop) & set(x)]
    else:
        xs_base = [x for x in xs_base if x not in xs_drop]
        return [x for x in xs if set(x) <= set(xs_base)]


def gen_xslist(xs: list[str] | list[list[str]],
               n_xs: int | range | None = None,
               xs_base: list[str] | None = None,
               xs_drop: list[str] | str | None = None) -> list[list[str]]:
    """ xs的组合，聚焦存在的指标清单，排除xs_drop

    Parameters
    ----------
    xs: list[str] | list[list[str]]
        待组合的指标清单，可以是字符串列表或指标组合列表
    n_xs: int | range | None, optional
        组合的长度，默认None，即所有长度的组合
    xs_base: list[str] | None, optional
        聚焦的指标清单，默认None，即所有指标
    xs_drop: list[str] | str | None, optional
        待排除的指标清单，可以是字符串或列表，默认None，即不排除任何指标

    Returns
    -------
    list[list[str]]
        指标组合的列表

    Examples
    --------
    >>> list_xs = ['a', 'b', 'c', 'd']
    >>> gen_xslist(list_xs, 2)
    [['a', 'b'], ['a', 'c'], ['a', 'd'], ['b', 'c'], ['b', 'd'], ['c', 'd']]
    >>> gen_xslist(list_xs, 2, xs_base=['a', 'b', 'c'])
    [['a', 'b'], ['a', 'c'], ['b', 'c']]
    >>> gen_xslist(list_xs, 2, xs_drop='a')
    [['b', 'c'], ['b', 'd'], ['c', 'd']]
    >>> gen_xslist(list_xs, 2, xs_base=['a', 'b'], xs_drop='a')
    [['b', 'c'], ['b', 'd']]
    >>> list_xs = [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]
    >>> gen_xslist(list_xs, 2)
    [['a', 'b'], ['c', 'd']]
    """
    if all(isinstance(x, str) for x in xs):
        return _gen_xs_wbd(xs, n_xs, xs_base, xs_drop)
    elif all(isinstance(x, list) for x in xs):
        return _gen_xslist_wbd(xs, n_xs, xs_base, xs_drop)
    else:
        raise ValueError('xs参数格式错误')


def gen_xs(xs: list[str],
           xs_drop: list[str] | None = None,
           xs_base: list[str] | None = None) -> list[str]:
    """ xs的组合，聚焦存在的指标清单，排除xs_drop """
    if xs_drop is None and xs_base is None:
        return xs
    elif xs_drop is not None and xs_base is None:
        return [x for x in xs if x not in xs_drop]
    elif xs_drop is None and xs_base is not None:
        return [x for x in xs if x in xs_base]
    else:
        return [x for x in xs if x in xs_base and xs not in xs_drop]
